#!/usr/bin/env python3
# -*- coding:utf-8 -*-

import torch
import torch.nn as nn

from .loss import build_criterion
from .yolov1 import YOLOv1

# YOLOv1模型的构建器，主要负责：
    # 构建YOLOv1模型实例
    # 初始化模型参数
    # 构建损失函数（如果需要训练）

# 构建 YOLOv1 网络
def build_yolov1(args, cfg, device, num_classes=80, trainable=False, deploy=False):
    print('==============================')
    print('Build {} ...'.format(args.model.upper()))
    
    print('==============================')
    print('Model Configuration: \n', cfg)
    
    # -------------- 构建YOLOv1 --------------
    model = YOLOv1(
        cfg = cfg,
        device = device,
        img_size = args.img_size, # 输入图像大小
        num_classes = num_classes, # 类别数量
        conf_thresh = args.conf_thresh, # 置信度阈值
        nms_thresh = args.nms_thresh, # NMS阈值
        trainable = trainable, # 是否是训练模式
        deploy = deploy
        )

    # -------------- 初始化YOLOv1的预测层参数 --------------
    # Init bias
    init_prob = 0.01 # 初始概率
    bias_value = -torch.log(torch.tensor((1. - init_prob) / init_prob)) # 计算偏置值

    # obj pred 目标存在性预测层参数初始化
    b = model.obj_pred.bias.view(1, -1)
    b.data.fill_(bias_value.item())
    model.obj_pred.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)

    # cls pred 类别预测层参数初始化
    b = model.cls_pred.bias.view(1, -1)
    b.data.fill_(bias_value.item())
    model.cls_pred.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)

    # reg pred 边界框预测层参数初始化
    b = model.reg_pred.bias.view(-1, )
    b.data.fill_(1.0)
    model.reg_pred.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) # 偏置初始化为1
    w = model.reg_pred.weight
    w.data.fill_(0.)
    model.reg_pred.weight = torch.nn.Parameter(w, requires_grad=True) # 权重初始化为0
    # 这样初始化可以确保初始化边界框预测接近网格中心

    # -------------- 构建用于计算标签分配和计算损失的Criterion类 --------------
    criterion = None
    if trainable:
        # build criterion for training 构建训练阶段使用的损失函数
        criterion = build_criterion(cfg, device, num_classes)

    return model, criterion

# 设计点
    # 参数初始化策略：
        # 使用特定的初始化值来确保模型训练初期的稳定性
        # 目标存在性和类别预测使用相同的初始化策略
        # 边界框预测使用不同的初始化策略
    # 模块化设计：
        # 将模型构建、参数初始化和损失函数构建分开
        # 便于代码维护和扩展
    # 条件构建：
        # 根据trainable参数决定是否构建损失函数
        # 根据deploy参数决定是否使用部署模式